



#ifndef AABB_TREE_H
#define AABB_TREE_H

#include <vector>
using std::vector;

#include <stack>
using std::stack;

#include "DrawUtils.h"

#include <OpenEXR/ImathBox.h> 
#include <OpenEXR/ImathBoxAlgo.h>
#include <OpenEXR/ImathVec.h>

using namespace Imath;

#include <limits.h>
#include <algorithm>
#include <assert.h>


template<typename T>
class AabbTree
{
	
public:
	
	AabbTree()
	{ }
	
	struct Node
	{
		Node() : m_left(INT_MAX), m_right(INT_MAX) {}
		
		int m_left;		 // left node index
		int m_right;	 // right node index
		Box3f m_aabb;	 // node AABB
		int m_leafIndex; // >=0 if and only if this is a leaf
		int m_skipNode;
		
		inline bool isLeaf() const { return m_leafIndex >= 0; }
	};
	
	struct TmpNode
	{
		TmpNode() : m_left(NULL), m_right(NULL) {}
		TmpNode* m_left;
		TmpNode* m_right;
		Box3f m_aabb; // node AABB
		int m_leafIndex;
		inline bool isLeaf() const { return m_leafIndex >= 0; }
	};
	
	Box3f getRootAabb() const;
	
	const Node& getNode(int nodeIndex) const { return m_nodes[nodeIndex]; }
	
	// Extract user data of a given leaf
	T getLeafData(int leafIndex) const;
	
	// Build tree top down, given the leaf AABBs and user infos.
	void buildTreeTopDown( const vector<Box3f>& leafAabbs, const vector<T>& leafDatas );
	
	// Update tree to accommodate new leaf AABBs.
	void refit( const vector<Box3f>& leafAabbs );
	
	// Given a test AABB, return the indices of the leaf nodes which it overlaps
	void query( const Box3f& testAabb, vector<int>& leafIndices ) const;
	
	// Empty the tree
	void clear();
	
	// Debug display
	void draw(int maxdepth, float r, float g, float b);
	
	friend class AabbTreeUtils;
	
public:
	
	struct Centroid
	{
		float m_centroid;
		int m_leafIndex;
	};
	
	inline friend bool operator<(const AabbTree<T>::Centroid& left, const AabbTree<T>::Centroid& right)
	{
		return left.m_centroid < right.m_centroid;
	}
	
private:
	
	vector<Node> m_nodes;
	vector<T> m_leafDatas;
	
	int m_numLeaves;
	int m_highestNode;
	
	TmpNode* m_rootTmpNode;
	
	void partitionObjects( const Box3f& nodeAabb, const vector<int>& leafIndices, vector<int>& leftSet, vector<int>& rightSet, const vector<Box3f>& leafAabbs );
	void buildTreeRecursive( TmpNode** node, const vector<int>& leafIndices, const vector<Box3f>& leafAabbs );
	void generateNodeIndicesRecursive(TmpNode* tmpnode, int i);
	void findSkipNodes();
	void freeTmpNodes(TmpNode* tmpnode);
	void _draw(int n, int depth, int maxdepth, float r, float g, float b);
	
};



class AabbTreeUtils
{
public:
	
	// Tree-tree collide
	struct NodePair 
	{
		// Node index m_a refers to tree A, Node index m_b refers to tree B
		NodePair(int a, int b) : m_a(a), m_b(b) {}
		int m_a, m_b;
	};
	
	struct SingleNode 
	{
		SingleNode(int index) : m_index(index) {}
		int m_index;
	};
	
	template<typename A, typename B>
	static void collideTrees( const AabbTree<A>& treeA, const AabbTree<B>& treeB, vector<NodePair>& leafPairs );
	
	template<typename T>
	static void selfCollide( const AabbTree<T>& tree, std::vector<NodePair>& leafPairs );
	
};



template<typename T>
void AabbTree<T>::buildTreeTopDown( const vector<Box3f>& leafAabbs, const vector<T>& leafDatas )
{
	assert(leafAabbs.size()>0);
	assert(leafAabbs.size() == leafDatas.size());
	
	m_numLeaves = leafAabbs.size();
	
	m_nodes.clear();
	m_leafDatas = leafDatas;
	
	vector<int> leafIndices;
	for (int i=0; i<leafDatas.size(); ++i)
	{
		leafIndices.push_back(i);
	}
	
	buildTreeRecursive(&m_rootTmpNode, leafIndices, leafAabbs);
	
	generateNodeIndicesRecursive(m_rootTmpNode, 0);
	
	findSkipNodes();
	
	freeTmpNodes(m_rootTmpNode);
}

template<typename T>
void AabbTree<T>::freeTmpNodes(TmpNode* tmpnode)
{
	if (tmpnode->m_left) freeTmpNodes(tmpnode->m_left);
	if (tmpnode->m_right) freeTmpNodes(tmpnode->m_right);
	delete tmpnode;
}

template<typename T>
void AabbTree<T>::clear()
{
	m_nodes.clear();
	m_leafDatas.clear();
	m_numLeaves = 0;
}

template<typename T>
void AabbTree<T>::generateNodeIndicesRecursive(TmpNode* tmpnode, int i)
{
	int numnodes = m_nodes.size();
	if (i>=numnodes) 
	{
		int expand = (i+1)*2;
		m_nodes.resize(expand);
	}
	
	Node& node = m_nodes[i];
	node.m_aabb = tmpnode->m_aabb;
	node.m_leafIndex = tmpnode->m_leafIndex;
	node.m_left = -1;
	node.m_right = -1;
	
	// Remap node pointers into an array of 2^(depth+1)-1 indices. Some of these elements will be dummies
	// (flagged by INT_MAX children). Since tree is nearly balanced due to use of median-cut heuristic, 
	// this is not too wasteful, and allows for very fast refitting (in which we skip the dummies).
	if (tmpnode->m_left)
	{
		m_nodes[i].m_left = 2*i+1;
		generateNodeIndicesRecursive(tmpnode->m_left, 2*i+1);
	}
	if (tmpnode->m_right)
	{
		m_nodes[i].m_right = 2*i+2;
		generateNodeIndicesRecursive(tmpnode->m_right, 2*i+2);
	}
}

template<typename T>
void AabbTree<T>::findSkipNodes()
{
	// Record in each non-dummy node the index of the previous non-dummy node (-1 flags the first).
	// This allows us to skip dummies when refitting, which is important because at high index there can be large numbers of contiguous dummies.
	int skipNode = -1;
	for (int n=0; n<m_nodes.size(); n++)
	{
		Node& node = m_nodes[n];
		bool isDummy = (node.m_left==INT_MAX);
		if (!isDummy)
		{
			node.m_skipNode = skipNode;
			skipNode = n;
		}
	}
	
	// Keep the index of the highest non-dummy node, to start from when refitting.
	m_highestNode = skipNode;
}

template<typename T>
void AabbTree<T>::buildTreeRecursive( TmpNode** node, const vector<int>& leafIndices, const vector<Box3f>& leafAabbs )
{
	assert(leafIndices.size() > 0);
	
	TmpNode* newNode = new TmpNode;
	*node = newNode;
	
	newNode->m_aabb = leafAabbs[leafIndices[0]];
	for(int i=1; i<leafIndices.size(); ++i)
	{
		newNode->m_aabb.extendBy(leafAabbs[leafIndices[i]]);
	}
	
	if (leafIndices.size() == 1)
	{
		newNode->m_left = NULL;
		newNode->m_right = NULL;
		newNode->m_leafIndex = leafIndices[0];
	}
	else
	{
		newNode->m_leafIndex = -1;
		
		// split the leaves indexed by leafIndices into two sets using the "median-cut" heuristic
		vector<int> leftSet;
		vector<int> rightSet;
		partitionObjects(newNode->m_aabb, leafIndices, leftSet, rightSet, leafAabbs);
		
		// Recursively construct left and right subtrees
		buildTreeRecursive(&(newNode->m_left), leftSet, leafAabbs);
		buildTreeRecursive(&(newNode->m_right), rightSet, leafAabbs);
	}
}


template<typename T>
void AabbTree<T>::partitionObjects( const Box3f& nodeAabb, const vector<int>& leafIndices, vector<int>& leftSet, vector<int>& rightSet, const vector<Box3f>& leafAabbs )
{
	// Choose partition axis as axis of maximum extent of nodeAabb
	V3f extents = nodeAabb.max - nodeAabb.min;
	int partitionAxis = nodeAabb.majorAxis();
	
	// Partition into two sets according to centroid location w.r.t. median centroid location.
	vector<Centroid> centroids;
	
	for (int i=0; i<leafIndices.size(); ++i)
	{
		const Box3f& aabb = leafAabbs[leafIndices[i]];
		V3f O = 0.5f * (aabb.max + aabb.min);
		
		Centroid C;
		C.m_centroid = O[partitionAxis];
		C.m_leafIndex = leafIndices[i];
		centroids.push_back(C);
	}
	
	sort(centroids.begin(), centroids.end());
	
	int median = centroids.size()/2;
	for (int i=0; i<median; ++i)
	{
		leftSet.push_back(centroids[i].m_leafIndex);
	}
	for (int i=median; i<centroids.size(); ++i)
	{
		rightSet.push_back(centroids[i].m_leafIndex);
	}
}


template<typename T>
void AabbTree<T>::draw(int maxdepth, float r, float g, float b)
{
	_draw(0, 0, maxdepth, r, g, b);
}

template<typename T>
void AabbTree<T>::_draw(int n, int depth, int maxdepth, float r, float g, float b)
{
	if (depth>maxdepth) return;
	
	Node& node = m_nodes[n];
	
	if (!node.isLeaf())
	{
		if (node.m_left>=0)
		{
			Node& leftNode = m_nodes[node.m_left];
			DrawUtils::drawBox(leftNode.m_aabb, r, g, b);
			_draw(node.m_left, depth+1, maxdepth, r, g, b);
		}
		if (node.m_right>=0)
		{
			Node& rightNode = m_nodes[node.m_right];
			DrawUtils::drawBox(rightNode.m_aabb, r, g, b);
			_draw(node.m_right, depth+1, maxdepth, r, g, b);
		}
	}
	else
	{
		Box3f aabb = node.m_aabb;
		DrawUtils::drawBox(0.5f*(aabb.max-aabb.min), r, g, b);
	}
}


template<typename T>
Box3f AabbTree<T>::getRootAabb() const
{
	const Node& root = m_nodes[0];
	return root.m_aabb;
}


template<typename T>
void AabbTree<T>::refit( const vector<Box3f>& leafAabbs )
{
	assert(leafAabbs.size()==m_numLeaves);
	
	// Since each child node has higher index than its parent node, we can just refit nodes in reverse order
	int currentNode = m_highestNode;
	
	while (currentNode > -1)
	{
		Node& node = m_nodes[currentNode];
		assert(node.m_left!=INT_MAX);
		
		if (!node.isLeaf())
		{
			// Not a leaf, compute AABB from children (both of which should have already been refitted)
			node.m_aabb.makeEmpty();
			
			assert(node.m_left>=0);
			assert(node.m_right>=0);
			
			Node& leftNode = m_nodes[node.m_left];
			node.m_aabb.extendBy(leftNode.m_aabb);
			
			Node& rightNode = m_nodes[node.m_right];
			node.m_aabb.extendBy(rightNode.m_aabb);
		}
		else
		{
			// A leaf, take AABB from provided new ones
			node.m_aabb = leafAabbs[node.m_leafIndex];
		}
		
		// Skip to the next non-dummy node.
		currentNode = node.m_skipNode;			
	}
}

template<typename T>
T AabbTree<T>::getLeafData(int leafIndex) const
{
	return m_leafDatas[leafIndex];
}


template<typename T>
void AabbTree<T>::query( const Box3f& testAabb, vector<int>& leafIndices ) const
{
	if (m_nodes.size()==0) return;
	
	stack<int> S;
	S.push(0);
	while (S.size())
	{
		int nodeIndex = S.top();
		S.pop();
		
		const Node& node = m_nodes[nodeIndex];
		if ( !testAabb.intersects(node.m_aabb) ) continue;
		
		if (node.isLeaf())
		{
			leafIndices.push_back(node.m_leafIndex);
		}
		else
		{
			assert(node.m_left>=0);
			assert(node.m_right>=0);
			S.push(node.m_left);
			S.push(node.m_right);
		}
	}
}


template<typename A, typename B>
static inline void _tandemTraverse( const AabbTree<A>& treeA, const AabbTree<B>& treeB,
								   int startNodeA, int startNodeB, std::vector<AabbTreeUtils::NodePair>& leafPairs )
{
	std::stack<AabbTreeUtils::NodePair> S;
	S.push(AabbTreeUtils::NodePair(startNodeA, startNodeB));
	
	while (S.size())
	{
		AabbTreeUtils::NodePair nodePair = S.top();
		S.pop();
		
		// Tandem-traverse nodeA and nodeB
		const class AabbTree<A>::Node& nodeA = treeA.getNode(nodePair.m_a);
		const class AabbTree<B>::Node& nodeB = treeB.getNode(nodePair.m_b);
		
		if ( !nodeA.m_aabb.intersects(nodeB.m_aabb) ) continue;
		
		if ( nodeA.isLeaf() && nodeB.isLeaf() )
		{
			AabbTreeUtils::NodePair leafPair(nodeA.m_leafIndex, nodeB.m_leafIndex);
			leafPairs.push_back(leafPair);
		}
		else
		{
			// Using "descend largest" heuristic. First determine node volumes.
			V3f extentsA = nodeA.m_aabb.size();
			V3f extentsB = nodeB.m_aabb.size();
			float volA = 1.f;
			float volB = 1.f;
			for (unsigned int d=0; d<3; ++d)
			{
				volA *= extentsA[d];
				volB *= extentsB[d];
			}
			
			// Descend A if largest and not a leaf (or we have to because B is a leaf)
			if ( nodeB.isLeaf() || (!nodeA.isLeaf() && (volA > volB)) )
			{
				S.push(AabbTreeUtils::NodePair(nodeA.m_left, nodePair.m_b));
				S.push(AabbTreeUtils::NodePair(nodeA.m_right, nodePair.m_b));
			}
			
			// Descend B
			else
			{
				S.push(AabbTreeUtils::NodePair(nodePair.m_a, nodeB.m_left));
				S.push(AabbTreeUtils::NodePair(nodePair.m_a, nodeB.m_right));
			}
		}
	}
}


template<typename A, typename B>
/* static */ void AabbTreeUtils::collideTrees( const AabbTree<A>& treeA, const AabbTree<B>& treeB, std::vector<NodePair>& leafPairs )
{
	_tandemTraverse<A,B>(treeA, treeB, 0, 0, leafPairs);
}


template<typename T>
/* static */ void AabbTreeUtils::selfCollide( const AabbTree<T>& tree, std::vector<NodePair>& leafPairs )
{
	std::stack<SingleNode> S;
	S.push(SingleNode(0));
	
	while (S.size())
	{
		SingleNode singleNode = S.top();
		S.pop();
		
		const class AabbTree<T>::Node& node = tree.getNode(singleNode.m_index);
		
		if ( node.isLeaf() ) continue;
		
		// Self-test left child
		S.push(SingleNode(node.m_left));
		
		// Self-test right child
		S.push(SingleNode(node.m_right));
		
		// Tandem-traverse children
		_tandemTraverse<T,T>(tree, tree, node.m_left, node.m_right, leafPairs);
	}
}


#endif //AABB_TREE_H


